#include "trajectory_generator_waypoint.h"
#include <stdio.h>
#include <ros/ros.h>
#include <ros/console.h>
#include <iostream>
#include <fstream>
#include <string>

using namespace std;    
using namespace Eigen;

TrajectoryGeneratorWaypoint::TrajectoryGeneratorWaypoint(){}
TrajectoryGeneratorWaypoint::~TrajectoryGeneratorWaypoint(){}

//define factorial function, input i, output i!
int TrajectoryGeneratorWaypoint::Factorial(int x)
{
    int fac = 1;
    for(int i = x; i > 0; i--)
        fac = fac * i;
    return fac;
}

Eigen::MatrixXd TrajectoryGeneratorWaypoint::getCt(int n_seg, int d_order){
    /*
    * get Selection Matrix
    * args:
    *   n_seg: the number of segments
    *   d_order: the differential order: if minimal jerk, d_oder=3
    * return:
    *   Ct: a matrix,
    *   row corresponds to original variable vector d ( 2 * d_order * n_seg )
    *   column corresponds to [ dF, dP ]' ( d_order*2*n_seg - (n_seg-1)*d_order )
    *   Note: the variables implicitly eliminates same variables""
    */
    int ct_rows = d_order * 2 * n_seg;
    int ct_cols = d_order * 2 * n_seg - (n_seg - 1) * d_order;
    MatrixXd Ct = MatrixXd::Zero(ct_rows, ct_cols);
    vector<int> d_vector;
    for (int k = 0; k < n_seg; ++k){
        for (int t = 0; t < 2; ++t){
            for (int d = 0; d < d_order; ++d){
                d_vector.push_back(100 * k + 10 * t + d);
            }
        }
    }
    int val, row;
    int col = 0;

    // fixed starting point at segment 0 in [dF, dP]'
    int k = 0;
    int t = 0;
    int d = 0;
    for (int d = 0; d < d_order; ++d){
        val = k * 100 + 10 * t + d;
        auto it = find(d_vector.begin(), d_vector.end(), val);
        row = distance(d_vector.begin(), it);
        Ct(row, col) = 1;
        col += 1;
    }

    // fixed final point at segment 0,1,2,...,n_seg-2 in [dF, dP]'
    k = 0;
    t = 1;
    d = 0;
    for (int k = 0; k < n_seg - 1; ++k){
        val = k * 100 + 10 * t + d;
        auto it = find(d_vector.begin(), d_vector.end(), val);
        row = distance(d_vector.begin(), it);
        Ct(row, col) = 1;

        val = (k + 1) * 100 + (t - 1) * 10 + d;
        it = find(d_vector.begin(), d_vector.end(), val);
        row = distance(d_vector.begin(), it);
        Ct(row, col) = 1;

        col += 1;
    }

    // fixed final point segment n_seg-1 in [dF, dP]'
    k = n_seg - 1;
    t = 1;
    d = 0;
    for (int d = 0; d < d_order; ++d){
        val = k * 100 + 10 * t + d;
        auto it = find(d_vector.begin(), d_vector.end(), val);
        row = distance(d_vector.begin(), it);
        Ct(row, col) = 1;
        col += 1;
    }

    // free variables at segment 0，1，2，n_seg-2 in [dF, dP]'
    k = 0;
    t = 1;
    d = 1;
    for (int k = 0; k < n_seg-1; ++k){
        for (int d = 1; d < d_order; ++d){
            val = k * 100 + 10 * t + d;
            auto it = find(d_vector.begin(), d_vector.end(), val);
            row = distance(d_vector.begin(), it);
            Ct(row, col) = 1;

            val = (k + 1) * 100 + (t - 1) * 10 + d;
            it = find(d_vector.begin(), d_vector.end(), val);
            row = distance(d_vector.begin(), it);
            Ct(row, col) = 1;

            col += 1;
        }
    }

    return Ct;
}

Eigen::MatrixXd TrajectoryGeneratorWaypoint::getM(int n_seg, int d_order, int p_num1d, const Eigen::VectorXd &ts){
    /*
     * get Mapping Matrix
     * args:
     *      n_seg: the number of segments
     *      d_order: the deferential order: if minimal jerk. it is 3,
     *      p_num1d: the number of variables in each segment, if minimal jerk, it is 6
     *      ts: time allocation as a column vector, size: n_seg x 1,
     * return:
     *      M: a matrix, size: n_seg * p_num1d x n_seg * p_num1d
     */
    cout << "in getM" << endl;
    MatrixXd M = MatrixXd::Zero(n_seg * p_num1d, n_seg * p_num1d);
    MatrixXd coeff(d_order, p_num1d);
    if(d_order == 4){
        coeff << 1, 1, 1, 1, 1, 1, 1, 1,
            0, 1, 2, 3, 4, 5, 6, 7,
            0, 0, 2, 6, 12, 20, 30, 42,
            0, 0, 0, 6, 24, 60, 120, 210;
    }
    else if(d_order == 3){
        coeff << 1, 1, 1, 1, 1, 1,
            0, 1, 2, 3, 4, 5,
            0, 0, 2, 6, 12, 20;
    }
    else{
        cout << "This derivatieve order is not provided getM!!!" << endl;
    }

    double t;
    for (int k = 0; k < n_seg; ++k){
        MatrixXd M_k = MatrixXd::Zero(p_num1d, p_num1d);
        t = ts(k);
        for (int i = 0; i < d_order; ++i){
            M_k(i, i) = coeff(i, i);
        }

        for (int i = 0; i < d_order; ++i){
            for (int j = i; j < p_num1d; ++j){
                M_k(i + d_order, j) = coeff(i, j) * pow(t, j - i);
            }
        }
        M.block(k * p_num1d, k * p_num1d, p_num1d, p_num1d) = M_k;
    }
    return M;
}

Eigen::MatrixXd TrajectoryGeneratorWaypoint::getQ(int n_seg, int d_order, int p_num1d, const Eigen::VectorXd &ts){
    /*
     * get Q matrix
     * args:
     *      n_seg: the number of segments
     *      d_order: the deferential order: if minimal snap. it is 4
     *      p_num1d: the number of variables in each segment
     *      ts: time allocation as a column vector, size: n_seg x 1,
     * return:
     *      Q: a matrix, size: n_seg * p_num1d x n_seg * p_num1d
     *      Note: p = [p0, p1, p2,...pn-1]'
     */
    MatrixXd Q = MatrixXd::Zero(n_seg * p_num1d, n_seg * p_num1d);
    for (int k = 0; k < n_seg; ++k){
        MatrixXd Q_k = MatrixXd::Zero(p_num1d, p_num1d);
        for (int i = 0; i < p_num1d; ++i){
            for (int j = 0; j < p_num1d; ++j){
                if(i<d_order || j<d_order)
                    continue;
                else
                    Q_k(i, j) = 1.0 * Factorial(i) * Factorial(j) / (Factorial(i - d_order) * Factorial(j - d_order) * (i + j - 2 * d_order + 1)) * pow(ts(k), i + j - 2 * d_order + 1);
            }
        }
        Q.block(k * p_num1d, k * p_num1d, p_num1d, p_num1d) = Q_k;
    }
    return Q;
}
/*

    STEP 2: Learn the "Closed-form solution to minimum snap" in L5, then finish this PolyQPGeneration function

    variable declaration: input       const int d_order,                    // the order of derivative
                                      const Eigen::MatrixXd &Path,          // waypoints coordinates (3d)
                                      const Eigen::MatrixXd &Vel,           // boundary velocity
                                      const Eigen::MatrixXd &Acc,           // boundary acceleration
                                      const Eigen::VectorXd &Time)          // time allocation in each segment
                          output      MatrixXd PolyCoeff(m, 3 * p_num1d);   // position(x,y,z), so we need (3 * p_num1d) coefficients

*/

Eigen::MatrixXd TrajectoryGeneratorWaypoint::PolyQPGeneration(
            const int d_order,                    // the order of derivative
            const Eigen::MatrixXd &Path,          // waypoints coordinates (3d)
            const Eigen::MatrixXd &Vel,           // boundary velocity
            const Eigen::MatrixXd &Acc,           // boundary acceleration
            const Eigen::VectorXd &Time)          // time allocation in each segment
{
    // enforce initial and final velocity and accleration, for higher order derivatives, just assume them be 0;
    int p_order   = 2 * d_order - 1;              // the order of polynomial
    int p_num1d   = p_order + 1;                  // the number of variables in each segment

    int n_seg = Time.size();                          // the number of segments
    MatrixXd PolyCoeff = MatrixXd::Zero(n_seg, 3 * p_num1d);           // position(x,y,z), so we need (3 * p_num1d) coefficients
    VectorXd Px(p_num1d * n_seg), Py(p_num1d * n_seg), Pz(p_num1d * n_seg);

    // get Q, get the cost function
    MatrixXd Q = getQ(n_seg, d_order, p_num1d, Time);
    cout << "get Q" << endl;

    /*   Produce Mapping Matrix M to the entire trajectory, M is a mapping matrix that maps polynomial coefficients to derivatives.   */
    Eigen::MatrixXd M = getM(n_seg, d_order, p_num1d, Time);
    cout << "get M" << endl;

    // Compute Ct
    MatrixXd Ct = getCt(n_seg, d_order);
    cout << "get Ct" << endl;

    MatrixXd C = Ct.transpose();

    MatrixXd M_inv = M.inverse();
    MatrixXd M_inv_t = M_inv.transpose();

    MatrixXd R = C * M_inv_t * Q * M_inv * Ct;

    int num_d_F = 2 * d_order  + n_seg - 1 ;
    int num_d_P = (n_seg - 1) * (d_order - 1);

    MatrixXd R_pp = R.bottomRightCorner(num_d_P, num_d_P);

    MatrixXd R_fp = R.topRightCorner(num_d_F, num_d_P);

    /*   Produce the dereivatives in X, Y and Z axis directly.  */
    for (int dim = 0; dim < 3; ++dim){
        VectorXd waypoints = Path.col(dim);

        VectorXd d_F = VectorXd::Zero(num_d_F);

        d_F(0) = waypoints(0); // p0 = 0
        // v0, a0, j0 = 0

        // PT,0, PT,1,  ,,PT,n_seg-2
        for (int i = 0; i < n_seg-1; ++i){
            d_F(d_order + i) = waypoints(i+1);
        }

        // PT,n_seg-1
        d_F(d_order + n_seg - 1) = waypoints(n_seg);

        VectorXd d_P = -1.0 * R_pp.inverse() * R_fp.transpose() * d_F;

        VectorXd d_total(d_F.rows() + d_P.rows());
        d_total << d_F, d_P;

        VectorXd poly_coef_1d = M_inv * Ct * d_total;

        MatrixXd poly_coef_1d_t = poly_coef_1d.transpose();

        for (int k = 0; k < n_seg; ++k){
            PolyCoeff.block(k, dim * p_num1d, 1, p_num1d) = poly_coef_1d_t.block(0, k * p_num1d, 1, p_num1d);
        }
    }
        /*   Produce the Minimum Snap cost function, the Hessian Matrix   */
        cout << "return PolyCoeff!!" << endl;
        return PolyCoeff;
        
}
